import json
import re

from utils.pre_process_question import process_question
from llm import llm, super_eval

prompt = '''
你将会回答用户的金融问题，主要为调用数据库查询工具，
数据库查询工具会基于你的问题进行解答
我会给你问题和预处理信息，你调用工具解答
调用格式为：
```json
{"query":"市场-问题"}
```

对于首个问题题，你只需要修改语病添加市场即可，无需改变。
示例：今天是2021年12月24日，创近半年新高的股票有几只？
回答：```json
{
    "query": "A股-2021年12月24日，创近半年新高的股票有几只？"
}
示例：2022年成立的CN公司有多少家？
回答：回答：```json
{
    "query": "美股-2022年成立的CN公司有多少家？"
}
需要补充完整一些属性定义,方便数据库查询工具匹配关键词。
示例：截止到目前，陕西省共有多少家上市公司？
回答：```json
{
    "query": "A股-截止到目前，公司注册地址为陕西省共有多少家上市公司？"
}
```
示例：华峰化学2019到2021的PBX值是多少？答案需要包含两位小数
回答：```json
{
    "query": "A股-华峰化学2019到2021的PBX（市净率的计算公式为：PBX = 总市值 / 股东权益合计)值是多少？答案需要包含两位小数"
}
```

如果有预处理信息，请将实体的类型，以及实体对应的所有代码，也输入到问题中，预处理信息会告诉你有哪些可选市场，问题只有一个市场，A股>港股>美股

示例
用户输入：
博时基金公司成立于（XXXX年XX月XX日）？{"success":true,"data":[{"ChiName":"博时基金管理有限公司","CompanyCode":41509}],"count":1}
你的输出：
```json
{
    "query": "A股-ChiName为博时基金管理有限公司（CompanyCode=41509）的公司成立于？输出格式：（XXXX年XX月XX日）"
}
```
如果查询得到答案：
博时基金公司成立于（1998年07月13日）
此时用户继续提问
比博时基金成立更早的基金公司(简称)有哪些？
则理解语义后继续输出：
```json
{
    "query": "A股-1998年07月13日之前成立的基金公司有哪些？回答基金简称。"
}
```

对于后面的问题你要理解语义并重写
示例2：
用户输入：
Q1:A股-在2021年全年，厦门钨业的最高收盘价是多少，出现在哪一天？{"InnerCode":1720,"CompanyCode":1485,"SecuCode":"600549","ChiName":"厦门钨业股份有限公司","ChiNameAbbr":"厦门钨业","EngName":"Xiamen Tungsten Co.,Ltd.","EngNameAbbr":"XTC","SecuAbbr":"厦门钨业","ChiSpelling":"XMWY"}
A1:【答案:32.93, 出现在2021年10月11日】
Q2:A股-2021年10月11日厦门钨业（InnerCode=1720，CompanyCode=1485）的VRSI量相对强弱是多少，保留4位小数？
A2:【答案:30.9743】
Q3:A股厦门钨业已连跌几天？
你的输出：
```json
{
    "query": "A股-2021年10月11日厦门钨业（InnerCode=1720，CompanyCode=1485）已连跌几天",
}
```

如果数据查询返回的内容可以回答，你则会按照```answer ```的格式进行回答，只回复关键词，关键词分开回复，关于名称类题目，优先考虑 `process_name` 字段结果。
例如：
问题:该股票是否创下了近一周，近一月或者近一季度新高？
回答方法：```answer
是;创下了近一周新高;创下了近一月新高;创下了近一季度新高
```

如果无法回答，则参考日志信息，看能否更加完善或者简化问题，进行再次提问。
注意你添加的code必须和实体匹配，慎重添加code，选择市场基于可选市场 A股>港股>美股，默认A股。
'''


class Controller:
    def __init__(self, solver):

        self.messages = [{"role": "system", "content": prompt}]
        self.solver = solver

    def run(self, question_info, is_first=False):
        question = question_info["question"]
        if is_first:
            market_dic = {'ConstantDB.SecuMain': "A股", 'ConstantDB.HK_SecuMain': "港股",
                          'ConstantDB.US_SecuMain': "美股"}

            info, tables = process_question(question)
            msg = ''
            if info:
                msg = f'\n系统消息：实体信息`{info}`'
            if tables:
                msg+='可选市场：'
                for table in tables:
                    msg += market_dic.get(table, '')

            self.messages.append({"role": "user", "content": question + f"\n{msg}"})

        else:

            self.messages.append(
                {"role": "user", "content": f"新问题为{question}，请基于历史信息调用工具，注意对实体添加实体类型，慎重添加code。"})

        solve_log = ''
        for try_num in range(5):
            answer = llm(self.messages)

            if '```json' in answer:
                query = super_eval(answer)['query']

                solve_log = self.solver.solve({"question":query,"id":question_info['id']})

                if try_num>1:
                    split_prompt = '请开始分析并决定输出答案或者继续查询，如果没有答案，这个问题可能比较难，请拆分成子问题一个一个解决。拆分方法为先查询过程实体ID然后进行下一步查询。'
                else:
                    split_prompt = '请开始分析并决定输出答案或者继续查询'

                self.messages.append({"role": "user", "content": f"""查询日志为：
--log--
{solve_log}
--end--
请基于日志分析:
原始问题是否被完全解答，如果问题被完全解答，请使用以下格式完整回答问题：
```answer
答案 
```
注意，答案不要修改原始答案的任何内容，数字中间不要加`,`，也不要擅自修改标点，除非有格式要求，否则日期和保留小数按照日志输出。
如果没有明确要求回答内容，对于名称类答案，统一回答`全称(简称)`的形式。优先查找`process_name`。
你需要考虑日志所有的内容，因为代码只会选择第一条作为答案，如果查询内容有多个可能的答案，你需要使用 `答案1;答案2`进行回复，其中答案2出现在日志但是没有被选择。
如果没有被完全解答，请将已知信息融入问题，并给出新的query。
原始问题为：{question}
{split_prompt}
"""})

            elif '```answer' in answer:
                final_answer = re.search('```answer(.*?)```', answer, re.DOTALL).group(1)
                # self.save_file(question_info, final_answer)
                return final_answer,solve_log
            else:
                self.messages.pop(-1)
        return '',''


    @staticmethod
    def save_file(question_info, final_answer):
        with open('submit.json','w', encoding='utf-8') as f:
            data = json.load(f)
        for i in data:
            for j in i['team']:
                if j['tid'] == question_info['tid']:
                    j['answer'] = final_answer
        with open('submit.json', 'w', encoding='utf-8') as f:
            f.write(json.dumps(data, ensure_ascii=False, indent=4))












